package com.hub.realtime.common.utils;

import lombok.extern.slf4j.Slf4j;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hadoop.yarn.api.records.ApplicationReport;
import org.apache.hadoop.yarn.api.records.YarnApplicationState;
import org.apache.hadoop.yarn.client.api.YarnClient;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.hadoop.yarn.exceptions.YarnException;

import java.io.IOException;

@Slf4j
public class YarnAdapter {
    private YarnClient yarnClient;
    public void init(String yarnConfDir) {
        Configuration conf = new YarnConfiguration();
        Configuration configuration = new YarnConfiguration(conf);
        configuration.addResource(new Path(yarnConfDir.concat("/yarn-site.xml")));
        yarnClient = YarnClient.createYarnClient();
        yarnClient.init(configuration);
        yarnClient.start();
    }
    public YarnApplicationState getApplicationState(String yarnConfDir,ApplicationId applicationId) {
        YarnApplicationState yarnApplicationState = null;
        try {
            init(yarnConfDir);
            ApplicationReport applicationReport = yarnClient.getApplicationReport(applicationId);
            yarnApplicationState = applicationReport.getYarnApplicationState();
        } catch (YarnException e) {
            log.error("client.getApplications YarnException", e);
        } catch (IOException e) {
            log.error("client.getApplications IOException", e);
        } finally {
            try {
                yarnClient.close();
            } catch (IOException e) {
                log.error("client.getApplications IOException", e);
            }
        }

        return yarnApplicationState;
    }

    public void killApplicationId(String yarnConfDir,ApplicationId applicationId) {
        try {
            init(yarnConfDir);
            yarnClient.killApplication(applicationId);
        } catch (YarnException e) {
            log.error("client.killApplicationId YarnException", e);
        } catch (IOException e) {
            log.error("client.killApplicationId IOException", e);
        } finally {
            try {
                yarnClient.close();
            } catch (IOException e) {
                log.error("client.getApplications IOException", e);
            }
        }

    }
}
